package com.lltt.study.controller;

import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.qdrant.QdrantEmbeddingStore;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.grpc.Collections;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Bean;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;

import static dev.langchain4j.store.embedding.filter.MetadataFilterBuilder.metadataKey;


@RestController
@Slf4j
public class EmbeddinglController {
    @Resource
    private EmbeddingModel embeddingModel;
    @Resource
    private QdrantClient qdrantClient;
    @Resource
    private EmbeddingStore<TextSegment> embeddingStore;

    @GetMapping(value = "/embedding/embed")
    public String embed(){
        String prompt = """
                   咏鸡
                鸡鸣破晓光，
                红冠映朝阳。
                金羽披霞彩，
                昂首步高岗。
                """;
        Response<Embedding> embeddingResponse = embeddingModel.embed(prompt);
        System.out.println("embeddingResponse = " + embeddingResponse);
        return embeddingResponse.content().toString();
    }
    //新建向量数据库实例和创建索引：test-qdrant
    @GetMapping(value = "/embedding/createCollection")
    public void createCollection(){
        var vectorParams = Collections.VectorParams.newBuilder()
                .setDistance(Collections.Distance.Cosine)
                .setSize(1024)
                .build();
        qdrantClient.createCollectionAsync("test-qdrant",vectorParams);
    }
    @GetMapping(value = "/embedding/add")
    public String add(){
        String prompt = """
                   咏鸡
                鸡鸣破晓光，
                红冠映朝阳。
                金羽披霞彩，
                昂首步高岗。
                """;
        TextSegment segment = TextSegment.from(prompt);
        segment.metadata().put("author","lltt");
        Embedding embedding = embeddingModel.embed(segment).content();
        String result = embeddingStore.add(embedding,segment);
        System.out.println(result);
        return result;
    }
    @GetMapping(value = "/embedding/query1")
    public void query1(){
        Embedding queryEmbedding  = embeddingModel.embed("咏鸡说的是什么").content();
        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)
                .maxResults(1)
                .build();
        EmbeddingSearchResult<TextSegment> search = embeddingStore.search(embeddingSearchRequest);
        System.out.println(search.matches().get(0).embedded().text());
    }
    @GetMapping(value = "/embedding/query2")
    public void query2(){
        Embedding queryEmbedding  = embeddingModel.embed("咏鸡").content();
        EmbeddingSearchRequest embeddingSearchRequest =
                EmbeddingSearchRequest.builder()
                        .queryEmbedding(queryEmbedding)
                        .filter(metadataKey("author").isEqualTo("lltt2"))
                        .maxResults(1)
                        .build();
        EmbeddingSearchResult<TextSegment> searchResult =
                embeddingStore.search(embeddingSearchRequest);
        System.out.println(searchResult.matches().get(0).embedded().text());
    }

}
